#!/usr/bin/env python
# coding: utf-8

# # 说明
#
# 我使用 Gemma-2 9b 4-bit 模型，**没有使用比赛之外的任何额外数据**，实现了 0.684 的 LB 分数。
#
# 为了避免破坏排行榜，我将模型参数保密，但我想分享我的实现，以证明仅使用 Gemma-2 9b 4-bit 模型即可实现此分数。
#
# 此 Notebook 改编自上一届比赛中 [@emiz6413](https://www.kaggle.com/emiz6413) 的最佳 Notebook 之一。
#
# 原始 Notebook 可以在[这里](https://www.kaggle.com/code/emiz6413/inference-gemma-2-9b-4-bit-qlora)找到。
#
# 我很高兴探索通过增加数据或切换到不同的模型，分数还能提高多少。
#
# 我希望这能对您有所帮助或启发！
#

# # 安装依赖库
# 使用 pip 安装所需的库，从本地文件加载，避免网络问题
import time
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor

import torch
import sklearn
import polars as pl
import numpy as np
import pandas as pd
from tqdm import tqdm
from transformers import (
    Gemma2ForSequenceClassification,
    GemmaTokenizerFast,
    BitsAndBytesConfig,
)
from transformers.data.data_collator import pad_without_fast_tokenizer_warning
from peft import PeftModel


# # 配置
@dataclass
class Config:
    # Gemma 2 9b 4bit 模型的路径
    gemma_dir = (
        "/kaggle/input/gemma-2/transformers/gemma-2-9b-it-4bit/1/gemma-2-9b-it-4bit"
    )
    # LoRA 模型的路径
    lora_dir = "/kaggle/input/wsdm2025-multilingual-chatbot-arena/open_model"
    # 最大输入长度
    max_length = 1900
    # 推理时使用的批次大小
    batch_size_list = [2, 16]
    # 是否进行测试时增强 (Test Time Augmentation, TTA)
    tta = True


cfg = Config()


# # 分词器
# 加载 Gemma 分词器
tokenizer = GemmaTokenizerFast.from_pretrained(cfg.gemma_dir)
# 添加 eos_token (end of sequence)
tokenizer.add_eos_token = True
# 设置填充方向为右侧
tokenizer.padding_side = "right"


# # 读取数据
# 读取测试数据
test_df = pl.read_parquet(
    "/kaggle/input/wsdm-cup-multilingual-chatbot-arena/test.parquet"
).to_pandas()


# # 预处理
# 对 prompt、response_a 和 response_b 列进行预处理，填充缺失值并处理过长文本
for col in ["prompt", "response_a", "response_b"]:
    test_df[col] = test_df[col].fillna("")  # 填充缺失值为空字符串
    text_list = []
    if col == "prompt":
        max_no = 402  # prompt 最大token数
        s_no = 200  # prompt 截断起始位置
        e_no = -201  # prompt 截断结束位置
    else:
        max_no = 702  # response 最大token数
        s_no = 350  # response 截断起始位置
        e_no = -351  # response 截断结束位置
    for text in tqdm(test_df[col]):  # 遍历每一段文本
        encoded = tokenizer(text, return_offsets_mapping=True)  # 分词，并返回偏移量
        if len(encoded["input_ids"]) > max_no:  # 如果文本长度超过最大值
            start_idx, end_idx = encoded["offset_mapping"][
                s_no
            ]  # 获取截断起始位置的字符偏移量
            new_text = text[:end_idx]  # 截取前半部分
            start_idx, end_idx = encoded["offset_mapping"][
                e_no
            ]  # 获取截断结束位置的字符偏移量
            new_text = (
                new_text + "\n(snip)\n" + text[start_idx:]
            )  # 将文本前后两部分拼接起来，中间用 (snip) 标记截断
            text = new_text  # 更新文本
        text_list.append(text)  # 将处理后的文本添加到列表中
    test_df[col] = text_list  # 更新 DataFrame 中的列


# # Tokenize
# 定义 tokenize 函数，将 prompt, response_a, response_b 拼接并分词
def tokenize(tokenizer, prompt, response_a, response_b, max_length=cfg.max_length):
    prompt = ["<prompt>: " + t for t in prompt]  # 在 prompt 前添加 "<prompt>: " 前缀
    response_a = [
        "\n\n<response_a>: " + t for t in response_a
    ]  # 在 response_a 前添加 "\n\n<response_a>: " 前缀
    response_b = [
        "\n\n<response_b>: " + t for t in response_b
    ]  # 在 response_b 前添加 "\n\n<response_b>: " 前缀
    texts = [
        p + r_a + r_b for p, r_a, r_b in zip(prompt, response_a, response_b)
    ]  # 将 prompt, response_a, response_b 拼接
    tokenized = tokenizer(
        texts, max_length=max_length, truncation=True
    )  # 使用分词器进行分词，并截断
    return (
        tokenized["input_ids"],
        tokenized["attention_mask"],
    )  # 返回 input_ids 和 attention_mask


# # 数据准备
# 对测试数据进行 tokenize，并准备增强数据，其中交换 response_a 和 response_b 的位置
data = pd.DataFrame()
data["id"] = test_df["id"]  # 创建 DataFrame，并添加 id 列
data["input_ids"], data["attention_mask"] = tokenize(
    tokenizer, test_df["prompt"], test_df["response_a"], test_df["response_b"]
)  # 使用 tokenize 函数进行分词
data["length"] = data["input_ids"].apply(len)  # 计算输入序列长度

aug_data = pd.DataFrame()
aug_data["id"] = test_df["id"]  # 创建增强数据 DataFrame，并添加 id 列
# 交换 response_a 和 response_b
aug_data["input_ids"], aug_data["attention_mask"] = tokenize(
    tokenizer, test_df["prompt"], test_df["response_b"], test_df["response_a"]
)  # 使用 tokenize 函数进行分词，交换 response_a 和 response_b
aug_data["length"] = aug_data["input_ids"].apply(len)  # 计算输入序列长度


# # 模型设置
# 在 GPU 0 上加载基础模型
device_0 = torch.device("cuda:0")
model_0 = Gemma2ForSequenceClassification.from_pretrained(
    cfg.gemma_dir,
    device_map=device_0,
    use_cache=False,  # 禁止使用缓存，减少显存占用
)
model_0.score = torch.nn.Linear(in_features=3584, out_features=2, bias=False).to(
    device_0
)  # 添加用于二分类的线性层

# 在 GPU 1 上加载基础模型
device_1 = torch.device("cuda:1")
model_1 = Gemma2ForSequenceClassification.from_pretrained(
    cfg.gemma_dir,
    device_map=device_1,
    use_cache=False,  # 禁止使用缓存，减少显存占用
)
model_1.score = torch.nn.Linear(in_features=3584, out_features=2, bias=False).to(
    device_1
)  # 添加用于二分类的线性层


# # 加载 LoRA 模型
# 加载 LoRA 模型，并将模型设置为评估模式
model_0 = PeftModel.from_pretrained(model_0, model_id=cfg.lora_dir).to(device_0)
model_0.eval()

model_1 = PeftModel.from_pretrained(model_1, model_id=cfg.lora_dir).to(device_1)
model_1.eval()


# # 推理函数
# 定义推理函数
@torch.no_grad()  # 关闭梯度计算
@torch.cuda.amp.autocast()  # 开启自动混合精度
def inference(df, model, device, batch_size, max_length=cfg.max_length):
    winners = []

    for start_idx in range(0, len(df), batch_size):  # 遍历每一个 batch
        end_idx = min(start_idx + batch_size, len(df))
        tmp = df.iloc[start_idx:end_idx]  # 获取当前 batch 的数据
        input_ids = tmp["input_ids"].to_list()  # 获取输入 id 列表
        attention_mask = tmp["attention_mask"].to_list()  # 获取 attention mask 列表
        inputs = pad_without_fast_tokenizer_warning(  # 将输入填充到相同长度
            tokenizer,
            {"input_ids": input_ids, "attention_mask": attention_mask},
            padding="longest",
            pad_to_multiple_of=None,
            return_tensors="pt",
        )
        outputs = model(**inputs.to(device))  # 模型推理
        proba = outputs.logits.softmax(-1).cpu()  # 对输出进行 softmax 并移到 CPU

        winners.extend(proba[:, 1].tolist())  # 将预测结果添加到 winners 列表中

    df["winner"] = winners  # 将预测结果添加到 DataFrame 中
    return df  # 返回 DataFrame


# # 执行推理
st = time.time()  # 记录推理起始时间
# 根据输入长度进行排序，以充分利用动态填充
if cfg.tta:  # 如果进行 TTA
    data = pd.concat([data, aug_data]).reset_index(drop=True)  # 合并原始数据和增强数据
    data["index"] = np.arange(len(data), dtype=np.int32)  # 添加索引列
    data = data.sort_values("length", ascending=False)  # 根据输入长度降序排序
else:
    data["index"] = np.arange(len(data), dtype=np.int32)  # 添加索引列
    data = data.sort_values("length", ascending=False)  # 根据输入长度降序排序

data_dict = {}
data_dict[0] = data[data["length"] > 1024].reset_index(
    drop=True
)  # 根据长度将数据分为两部分
data_dict[1] = data[data["length"] <= 1024].reset_index(
    drop=True
)  # 根据长度将数据分为两部分
result_df = []
for i, batch_size in enumerate(Config.batch_size_list):
    if len(data_dict[i]) == 0:
        continue
    sub_1 = data_dict[i].iloc[0::2].copy()  # 将数据分为两部分
    sub_2 = data_dict[i].iloc[1::2].copy()  # 将数据分为两部分

    with ThreadPoolExecutor(max_workers=2) as executor:  # 使用多线程进行推理
        results = executor.map(
            inference,
            (sub_1, sub_2),
            (model_0, model_1),
            (device_0, device_1),
            (batch_size, batch_size),
        )  # 执行推理

    result_df.append(pd.concat(list(results), axis=0))  # 将推理结果合并

result_df = (
    pd.concat(result_df).sort_values("index").reset_index(drop=True)
)  # 合并所有结果并按索引排序
if cfg.tta:  # 如果进行 TTA
    proba = result_df["winner"].values[: len(aug_data)]  # 获取原始数据的预测结果
    tta_proba = (
        1 - result_df["winner"].values[len(aug_data) :]
    )  # 获取增强数据的预测结果
    proba = (proba + tta_proba) / 2  # 将原始数据和增强数据的预测结果取平均
    result_df = result_df[: len(aug_data)].copy()  # 只保留原始数据的预测结果
    result_df["winner"] = proba  # 更新预测结果
print(f"elapsed time: {time.time() - st}")  # 打印推理时间


# # 生成提交文件
submission_df = result_df[["id", "winner"]].copy()  # 创建提交 DataFrame
submission_df["winner"] = np.where(
    submission_df["winner"] > 0.5, "model_a", "model_b"
)  # 根据预测概率选择模型
submission_df.to_csv("submission.csv", index=False)  # 保存提交文件
display(result_df[["id", "winner"]])  # 显示预测结果
